package codegen


import . "compiler/ast"


func cgStat(fi *funcInfo, node Stat) {

	switch stat := node.(type) {
		case *FuncCallStat:		cgFuncCallStat(fi, stat)

		case *BreakStat:		cgBreakStat(fi, stat)

		case *DoStat:			cgDoStat(fi, stat)

		case *RepeatStat:		cgRepeatStat(fi, stat)

		case *WhileStat:		cgWhileStat(fi, stat)

		case *IfStat:			cgIfStat(fi, stat)

		case *ForNumStat:		cgForNumStat(fi, stat)

		case *ForInStat:		cgForInStat(fi, stat)

		case *AssignStat:		cgAssignStat(fi, stat)

		case *LocalVarDeclStat:	cgLocalVarDeclStat(fi, stat)

		case *LocalFuncDefStat:	cgLocalFuncDefStat(fi, stat)

		case *LabelStat, *GotoStat:		panic("label and goto statements are not supported!")
	}
}


// 局部函数定义语句
// local function f() end完全等价于local f; f = function() end
func cgLocalFuncDefStat(fi *funcInfo, node *LocalFuncDefStat) {

	r := fi.addLocVar(node.Name)
	cgFuncDefExp(fi, node.Exp, r)
}


// 对于函数调用语句，可以认为是对函数调用表达式进行求值但不需要任何返回值
func cgFuncCallStat(fi *funcInfo, node *FuncCallStat) {

	r := fi.allocReg()
	cgFuncCallExp(fi, node, r, 0)

	fi.freeReg()
}


// 对于break语句，生成一条JMP指令并把地址记录到break表里即可
// 等到块退出以后再修补跳转偏移量
func cgBreakStat(fi *funcInfo, node *BreakStat) {

	pc := fi.emitJmp(0, 0)
	fi.addBreakJump(pc)

}


// 由于do语句本质上就是一个块，其用途是引入新的作用域
func cgDoStat(fi *funcInfo, node *DoStat) {

	fi.enterScope(false)		// 非循环块
	cgBlock(fi, node.Block)
	fi.closeOpenUpvals()
	fi.exitScope()
}


/*
           ______________
          /  false? jmp  |
         /               |
while exp do block end <-'
      ^           \
      |___________/
           jmp
*/ 
// 表达式为: while a == 1 do local e = 1 end 
// 对于这种代码生成方式，生成的code相比于Lua官方源码
// 会多一条TEST指令
// 但执行结果是一样的	    
func cgWhileStat(fi *funcInfo, node *WhileStat) {

	// 首先记住当前的pc，因为后面计算跳转偏移量是时要用到
	pcBeforeExp := fi.pc()

	// 分配一个临时变量，对表达式进行求值，然后释放临时变量
	r := fi.allocReg()
	cgExp(fi, node.Exp, r, 1)
	fi.freeReg()

	// 生成TEST和JMP指令，实现条件跳转
	// 由于此时还没有对块进行处理，所以跳转的偏移量还没办法给出
	fi.emitTest(r, 0)
	pcJumpToEnd := fi.emitJmp(0, 0)

	// 对块进行处理，生成一条JMP指令跳转到最开始
	fi.enterScope(true)
	cgBlock(fi, node.Block)
	fi.closeOpenUpvals()
	fi.emitJmp(0, pcBeforeExp-fi.pc()-1)  // -1的原因是pc会自动+1
	fi.exitScope()

	// 修复第一条JMP指令的偏移量
	fi.fixSbx(pcJumpToEnd, fi.pc()-pcJumpToEnd)  // 跳转到while语句最后一条指令的下一条指令，由于会自动+1，因此不需要+1
}


/*
        ______________
       |  false? jmp  |
       V              /
repeat block until exp
*/
func cgRepeatStat(fi *funcInfo, node *RepeatStat) {

	fi.enterScope(true)

	pcBeforeBlock := fi.pc()
	cgBlock(fi, node.Block)

	r := fi.allocReg()
	cgExp(fi, node.Exp, r, 1)
	fi.freeReg()

	fi.emitTest(r, 0)
	fi.emitJmp(fi.getJmpArgA(), pcBeforeBlock-fi.pc()-1)
	fi.closeOpenUpvals()

	// repeat语句块的作用域时把后面的表达式也覆盖在内的
	// 所以在表达式里可以访问到块里声明的局部变量
	fi.exitScope()
}


/*
         _________________       _________________       _____________
        / false? jmp      |     / false? jmp      |     / false? jmp  |
       /                  V    /                  V    /              V
if exp1 then block1 elseif exp2 then block2 elseif true then block3 end <-.
                   \                       \                       \      |
                    \_______________________\_______________________\_____|
                    jmp                     jmp                     jmp
*/
func cgIfStat(fi *funcInfo, node *IfStat) {

	pcJmpToEnds := make([]int, len(node.Exps))
	pcJmpToNextExp := -1

	for i, exp := range node.Exps {
		// 将前len-1个分支的exp语句的跳转指向下一个分支的exp语句的pc
		if pcJmpToNextExp >= 0 {
			fi.fixSbx(pcJmpToNextExp, fi.pc()-pcJmpToNextExp)
		}

		r := fi.allocReg()
		cgExp(fi, exp, r, 1)
		fi.freeReg()

		fi.emitTest(r, 0)
		pcJmpToNextExp = fi.emitJmp(0, 0)

		fi.enterScope(false)
		cgBlock(fi, node.Blocks[i])
		fi.closeOpenUpvals()
		fi.exitScope()

		// 最后一个条件分支，block结尾不需要增加jmp指令
		// 直接pc+1就到整个if语句的下一条指令
		// 但是对于除最后一个以外的分支，都需要增加跳转指令
		// 在对应的block执行结束时，可以调到整个if语句的下一条指令
		if i < len(node.Exps) - 1 {
			pcJmpToEnds[i] = fi.emitJmp(0, 0)
		} else {
			pcJmpToEnds[i] = pcJmpToNextExp
		}
	}

	for _, pc := range pcJmpToEnds {
		fi.fixSbx(pc, fi.pc() - pc)
	}
}


// 数值for循环
func cgForNumStat(fi *funcInfo, node *ForNumStat) {

	fi.enterScope(true)		// 循环

	// 数值for循环使用了三个特殊的局部变量
	// 分别存储索引、限制和步长
	// 所以第一步就是声明这三个局部变量
	// 并且使用初始、限制和步长表达式的值初始化这三个变量
	cgLocalVarDeclStat(fi, &LocalVarDeclStat{
		NameList:	[]string{"(for index)", "(for limit)", "(for step)"},
		ExpList:	[]Exp{node.InitExp, node.LimitExp, node.StepExp},
	})
	// 关键字for后面的名字也声明了一个局部变量
	fi.addLocVar(node.VarName)

	// 生成FORPREP指令，处理块，然后生成FORLOOP指令
	a := fi.usedRegs - 4	// usedRegs指向当前可用的寄存器，for声明一共占了4个变量，-4目的是获得for语句开始的变量寄存器地址
	pcForPrep := fi.emitForPrep(a, 0)
	cgBlock(fi, node.Block)
	fi.closeOpenUpvals()
	pcForLoop := fi.emitForLoop(a, 0)

	// 把指令里的调整偏移量修复即可
	fi.fixSbx(pcForPrep, pcForLoop-pcForPrep-1)  // pcForPrep跳转到pcForPrep
	fi.fixSbx(pcForLoop, pcForPrep-pcForLoop)	// ForPrep仅在开始执行一次, 所以这里加上pc会自动+1, 其实是跳到pcForPrep的下一条语句

	fi.exitScope()
} 


// 通用for循环
func cgForInStat(fi *funcInfo, node *ForInStat) {

	fi.enterScope(true)

	cgLocalVarDeclStat(fi, &LocalVarDeclStat{
		NameList:	[]string{"(for generator)", "(for state)", "(for control)"},
		ExpList:	node.ExpList,
	})
	for _, name := range node.NameList {
		fi.addLocVar(name)
	}

	pcJmpToTFC := fi.emitJmp(0, 0)
	cgBlock(fi, node.Block)
	fi.closeOpenUpvals()
	fi.fixSbx(pcJmpToTFC, fi.pc()-pcJmpToTFC)

	rGenerator := fi.slotOfLocVar("(for generator)")
	fi.emitTForCall(rGenerator, len(node.NameList))
	fi.emitTForLoop(rGenerator+2, pcJmpToTFC-fi.pc()-1)  // fi.pc()指向的还是上一条指令的地址, 为了跳转到pcJmpToTFC的下一条指令, 需要再-1

	fi.exitScope()

}


// 局部局部变量声明
// 由于局部变量声明语句可以一次声明多个局部变量，并对变量进行初始化
func cgLocalVarDeclStat(fi *funcInfo, node *LocalVarDeclStat) {
	exps := removeTailNils(node.ExpList)  // 移除表达式末尾的nil
	nExps := len(exps)
	nNames := len(node.NameList)

	oldRegs := fi.usedRegs

	// 如果等号左侧声明的局部变量和等号右侧提供的表达式一样多
	if nExps == nNames {
		// 新声明的局部变量只有在声明语句结束之后才生效  // TODO ?
		// 所以先分配临时内存，对表达式求值，后续再将局部变量名和寄存器绑定
		for _, exp := range exps {
			a := fi.allocReg()
			cgExp(fi, exp, a, 1)
		}
	} else if nExps > nNames { 	// 表达式比局部变量多
		// 多余的表达式一样要求值
		// 如果最后一个表达式是vararg或者函数调用
		// 还需要特别处理一下
		for i, exp := range exps {
			a := fi.allocReg()
			if i == nExps - 1 && isVarargOrFuncCall(exp) {
				cgExp(fi, exp, a, 0)
			} else {
				cgExp(fi, exp, a, 1)
			}
		}
	} else {  // 表达式比局部变量少
		multRet := false
		for i, exp := range exps {
			a := fi.allocReg()
			// 如果最后一个表达式是vararg或者函数调用
			// 则需要使用多重赋值初始化其余的变量
			// 否则，必须生成LOADNIL指令来初始化剩余的局部变量
			if i == nExps - 1 && isVarargOrFuncCall(exp) {
				multRet = true
				n := nNames - nExps + 1
				cgExp(fi, exp, a, n)
				fi.allocRegs(n-1)
			} else {
				cgExp(fi, exp, a, 1)
			}
		}

		if !multRet {
			n := nNames - nExps
			a := fi.allocRegs(n)
			fi.emitLoadNil(a, n)
		}
	}

	// 释放临时变量，声明局部变量
	fi.usedRegs = oldRegs
	for _, name := range node.NameList {
		fi.addLocVar(name)
	}
}


// 赋值语句不仅支持多重赋值
// 而且还可以同时给局部变量、Upvalue、全局变量赋值或者根据索引修改表
func cgAssignStat(fi *funcInfo, node *AssignStat) {
	exps := removeTailNils(node.ExpList)	// 移除表达式末尾的nil
	nExps := len(exps)
	nVars := len(node.VarList)

	oldRegs := fi.usedRegs

	// 由于赋值语句等号左边可以出现t[k]这样的表达式
	// 等号右边则可以出现任意表达式，所以我们需要先分配临时变量
	// 对这些表达式进行求值
	// 然后在统一生产赋值指令
	// tRegs、kRegs和vRegs这三个数组分别记录我们表、键和值分配的临时变量
	tRegs := make([]int, nVars)
	kRegs := make([]int, nVars)
	vRegs := make([]int, nVars)

	// 先处理等号左侧的索引表达式，分配临时变量
	// 并对表和键求值
	for i, exp := range node.VarList {
		// 在_finishVarList中通过_checkVar将var的类型限定为NameExp和TableAccessExp
		if taExp, ok := exp.(*TableAccessExp); ok {
			tRegs[i] = fi.allocReg()
			cgExp(fi, taExp.PrefixExp, tRegs[i], 1)
			kRegs[i] = fi.allocReg()
			cgExp(fi, taExp.KeyExp, kRegs[i], 1)
		} else {
			name := exp.(*NameExp).Name
			if fi.slotOfLocVar(name) < 0 && fi.indexOfUpval(name) < 0 {
				// global var
				kRegs[i] = -1  // TODO ?
				if fi.indexOfConstant(name) > 0xFF {
					kRegs[i] = fi.allocReg()
				}
			}
		}
	}

	// 统一为等号右侧的表达式计算寄存器索引
	for i := 0; i < nVars; i++ {
		vRegs[i] = fi.usedRegs + i
	}

	if nExps >= nVars {
		for i, exp := range exps {
			a := fi.allocReg()
			if i >= nVars && i == nExps-1 && isVarargOrFuncCall(exp) {  // TODO ???
				cgExp(fi, exp, a, 0)  // ...或者函数调用的结果放在栈上? TODO ?
			} else {
				cgExp(fi, exp, a, 1)
			}
		}
	} else {
		multRet := false

		for i, exp := range exps {
			a := fi.allocReg()
			if i == nExps-1 && isVarargOrFuncCall(exp) {
				multRet = true
				n := nVars - nExps + 1
				cgExp(fi, exp, a, n)  // ...或者函数调用的结果取出n+1个赋值给变量? TODO?
				fi.allocRegs(n-1)
			} else {
				cgExp(fi, exp, a, 1)
			}
		}

		// 考虑多重赋值，需要在必要时补上LOADNIL指令
		if !multRet {
			n := nVars - nExps
			a := fi.allocRegs(n)
			fi.emitLoadNil(a, n)
		}
	}

	// 在循环中对赋值进行处理
	// 如果给局部变量赋值，生成MOVE指令
	// 如果给Upvalue赋值，生成SETUPVAL指令
	// 如果给全局变量赋值，生成SETTABUP指令
	// 如果是按索引给表赋值，则生成SETTABLE指令
	for i, exp := range node.VarList {
		if nameExp, ok := exp.(*NameExp); ok {
			varName := nameExp.Name
			if a := fi.slotOfLocVar(varName); a >= 0 {
				fi.emitMove(a, vRegs[i])
			} else if b := fi.indexOfUpval(varName); b >= 0 {
				fi.emitSetUpval(vRegs[i], b)
			} else if a := fi.slotOfLocVar("_ENV"); a >= 0 {  // _ENV是main函数的一个变量
				if kRegs[i] < 0 {
					b := 0x100 + fi.indexOfConstant(varName)
					fi.emitSetTable(a, b, vRegs[i])
				} else {
					fi.emitSetTable(a, kRegs[i], vRegs[i])
				}
			} else {		// global var
				a := fi.indexOfUpval("_ENV")
				if kRegs[i] < 0 {
					b := 0x100 + fi.indexOfConstant(varName)
					fi.emitSetTabUp(a, b, vRegs[i])
				} else {
					fi.emitSetTabUp(a, kRegs[i], vRegs[i])
				}
			}
		} else {
			fi.emitSetTable(tRegs[i], kRegs[i], vRegs[i])
		}
	}

	// 循环之后释放所有临时变量
	fi.usedRegs = oldRegs
}